"""
数据处理工具函数文件

本文件提供脑波数据处理系统中常用的数据转换、格式化和处理工具函数。
主要用于支持模块间数据传递时的格式转换和数据处理需求。

使用方法：
- create_time_series_data(): 创建时域数据对象
- create_frequency_data(): 创建频域数据对象  
- create_brain_wave_data(): 创建完整脑波数据对象
- convert_raw_to_time_series(): 原始数据转时域格式
- merge_brain_wave_data(): 合并多个数据对象

原理：
提供标准化的数据创建和转换接口，简化各模块的数据处理工作。
支持不同数据源的格式统一，确保数据流的一致性。
"""

import numpy as np
from datetime import datetime
from typing import List, Dict, Any, Optional, Union
from ..models.brain_wave_data import (
    BrainWaveData, 
    TimeSeriesData, 
    FrequencyData, 
    DataValidator,
    DataConstants
)


class DataUtils:
    """数据处理工具类"""
    
    @staticmethod
    def create_time_series_data(
        data: Union[List[float], np.ndarray],
        sampling_rate: int = DataConstants.TIME_SERIES_SAMPLING_RATE,
        channel_names: Optional[List[str]] = None,
        timestamp: Optional[datetime] = None
    ) -> TimeSeriesData:
        """
        创建时域数据对象
        
        Args:
            data: 原始数据，可以是列表或numpy数组
            sampling_rate: 采样率，默认512Hz
            channel_names: 通道名称列表
            timestamp: 时间戳
            
        Returns:
            TimeSeriesData: 时域数据对象
        """
        # 转换数据格式
        if isinstance(data, list):
            data = np.array(data, dtype=np.float32)
        elif not isinstance(data, np.ndarray):
            raise ValueError("数据必须是列表或numpy数组")
        
        # 设置默认值
        if channel_names is None:
            if data.ndim == 1:
                channel_names = ['EEG']
            else:
                channel_names = [f'EEG_{i+1}' for i in range(data.shape[0])]
        
        if timestamp is None:
            timestamp = datetime.now()
        
        return TimeSeriesData(
            data=data,
            sampling_rate=sampling_rate,
            timestamp=timestamp,
            channel_names=channel_names
        )
    
    @staticmethod
    def create_frequency_data(
        # 8个频段功率数据
        delta: float = 0.0,
        theta: float = 0.0,
        low_alpha: float = 0.0,
        high_alpha: float = 0.0,
        low_beta: float = 0.0,
        high_beta: float = 0.0,
        low_gamma: float = 0.0,
        mid_gamma: float = 0.0,
        
        # NeuroSky算法输出
        attention: float = 0.0,
        meditation: float = 0.0,
        
        # 信号质量和其他指标
        poor_signal: int = 0,
        blink_strength: int = 0,
        
        # 原始数据
        raw_value: int = 0,
        
        # 元数据
        is_filtered: bool = False,
        timestamp: Optional[datetime] = None
    ) -> FrequencyData:
        """
        创建频域数据对象 - 基于NeuroPy3的完整数据字段
        
        Args:
            # 8个频段功率数据
            delta: Delta波段功率
            theta: Theta波段功率
            low_alpha: 低Alpha波段功率
            high_alpha: 高Alpha波段功率
            low_beta: 低Beta波段功率
            high_beta: 高Beta波段功率
            low_gamma: 低Gamma波段功率
            mid_gamma: 中Gamma波段功率
            
            # NeuroSky算法输出
            attention: 注意力指数
            meditation: 冥想指数
            
            # 信号质量和其他指标
            poor_signal: 信号质量指标
            blink_strength: 眨眼强度
            
            # 原始数据
            raw_value: 原始EEG数据值
            
            # 元数据
            is_filtered: 是否为过滤后的数据
            timestamp: 时间戳
            
        Returns:
            FrequencyData: 频域数据对象
        """
        if timestamp is None:
            timestamp = datetime.now()
        
        return FrequencyData(
            # 8个频段功率数据
            delta=delta,
            theta=theta,
            low_alpha=low_alpha,
            high_alpha=high_alpha,
            low_beta=low_beta,
            high_beta=high_beta,
            low_gamma=low_gamma,
            mid_gamma=mid_gamma,
            
            # NeuroSky算法输出
            attention=attention,
            meditation=meditation,
            
            # 信号质量和其他指标
            poor_signal=poor_signal,
            blink_strength=blink_strength,
            
            # 原始数据
            raw_value=raw_value,
            
            # 元数据
            is_filtered=is_filtered,
            timestamp=timestamp
        )
    
    @staticmethod
    def create_brain_wave_data(
        time_series: TimeSeriesData,
        frequency_data: Optional[FrequencyData] = None,
        session_id: str = "",
        processing_stage: str = "raw",
        metadata: Optional[Dict[str, Any]] = None
    ) -> BrainWaveData:
        """
        创建完整脑波数据对象
        
        Args:
            time_series: 时域数据
            frequency_data: 频域数据（可选）
            session_id: 会话ID
            processing_stage: 处理阶段
            metadata: 元数据
            
        Returns:
            BrainWaveData: 完整脑波数据对象
        """
        if metadata is None:
            metadata = {}
        
        return BrainWaveData(
            time_series=time_series,
            frequency_data=frequency_data,
            session_id=session_id,
            processing_stage=processing_stage,
            metadata=metadata
        )
    
    @staticmethod
    def convert_raw_to_time_series(
        raw_data: Union[List[float], np.ndarray],
        sampling_rate: int = DataConstants.TIME_SERIES_SAMPLING_RATE,
        session_id: str = ""
    ) -> BrainWaveData:
        """
        将原始数据转换为标准时域格式
        
        Args:
            raw_data: 原始数据
            sampling_rate: 采样率
            session_id: 会话ID
            
        Returns:
            BrainWaveData: 包含时域数据的脑波数据对象
        """
        time_series = DataUtils.create_time_series_data(
            data=raw_data,
            sampling_rate=sampling_rate
        )
        
        return DataUtils.create_brain_wave_data(
            time_series=time_series,
            session_id=session_id,
            processing_stage="raw"
        )
    
    @staticmethod
    def add_frequency_data_to_brain_wave(
        brain_wave_data: BrainWaveData,
        frequency_data: FrequencyData
    ) -> BrainWaveData:
        """
        向脑波数据对象添加频域数据
        
        Args:
            brain_wave_data: 原始脑波数据对象
            frequency_data: 频域数据
            
        Returns:
            BrainWaveData: 更新后的脑波数据对象
        """
        brain_wave_data.frequency_data = frequency_data
        return brain_wave_data
    
    @staticmethod
    def update_processing_stage(
        brain_wave_data: BrainWaveData,
        new_stage: str
    ) -> BrainWaveData:
        """
        更新数据处理阶段标识
        
        Args:
            brain_wave_data: 脑波数据对象
            new_stage: 新的处理阶段
            
        Returns:
            BrainWaveData: 更新后的脑波数据对象
        """
        if new_stage not in DataConstants.PROCESSING_STAGES:
            raise ValueError(f"无效的处理阶段: {new_stage}")
        
        brain_wave_data.processing_stage = new_stage
        return brain_wave_data
    
    @staticmethod
    def merge_time_series_data(
        data_list: List[TimeSeriesData]
    ) -> TimeSeriesData:
        """
        合并多个时域数据对象
        
        Args:
            data_list: 时域数据对象列表
            
        Returns:
            TimeSeriesData: 合并后的时域数据
        """
        if not data_list:
            raise ValueError("数据列表不能为空")
        
        # 检查采样率一致性
        sampling_rate = data_list[0].sampling_rate
        if not all(data.sampling_rate == sampling_rate for data in data_list):
            raise ValueError("所有数据的采样率必须一致")
        
        # 合并数据
        if data_list[0].data.ndim == 1:
            merged_data = np.concatenate([data.data for data in data_list])
        else:
            merged_data = np.concatenate([data.data for data in data_list], axis=1)
        
        # 使用第一个数据的元信息
        return TimeSeriesData(
            data=merged_data,
            sampling_rate=sampling_rate,
            timestamp=data_list[0].timestamp,
            channel_names=data_list[0].channel_names
        )
    
    @staticmethod
    def extract_data_window(
        brain_wave_data: BrainWaveData,
        start_time: float,
        duration: float
    ) -> BrainWaveData:
        """
        从脑波数据中提取指定时间窗口的数据
        
        Args:
            brain_wave_data: 原始脑波数据
            start_time: 开始时间（秒）
            duration: 持续时间（秒）
            
        Returns:
            BrainWaveData: 提取的数据窗口
        """
        sampling_rate = brain_wave_data.time_series.sampling_rate
        start_sample = int(start_time * sampling_rate)
        end_sample = int((start_time + duration) * sampling_rate)
        
        # 提取时域数据窗口
        if brain_wave_data.time_series.data.ndim == 1:
            windowed_data = brain_wave_data.time_series.data[start_sample:end_sample]
        else:
            windowed_data = brain_wave_data.time_series.data[:, start_sample:end_sample]
        
        # 创建新的时域数据对象
        windowed_time_series = TimeSeriesData(
            data=windowed_data,
            sampling_rate=sampling_rate,
            timestamp=brain_wave_data.time_series.timestamp,
            channel_names=brain_wave_data.time_series.channel_names
        )
        
        # 创建新的脑波数据对象
        return BrainWaveData(
            time_series=windowed_time_series,
            frequency_data=brain_wave_data.frequency_data,  # 保持原有频域数据
            session_id=brain_wave_data.session_id,
            processing_stage=brain_wave_data.processing_stage,
            metadata=brain_wave_data.metadata.copy()
        )
    
    @staticmethod
    def validate_and_clean_data(
        brain_wave_data: BrainWaveData,
        remove_invalid: bool = True
    ) -> BrainWaveData:
        """
        验证并清理数据
        
        Args:
            brain_wave_data: 脑波数据对象
            remove_invalid: 是否移除无效数据
            
        Returns:
            BrainWaveData: 清理后的数据对象
        """
        # 验证数据
        if not DataValidator.validate_brain_wave_data(brain_wave_data):
            if not remove_invalid:
                raise ValueError("数据验证失败")
        
        # 清理时域数据中的异常值
        time_data = brain_wave_data.time_series.data.copy()
        
        # 移除NaN和无穷大值
        if remove_invalid:
            if time_data.ndim == 1:
                valid_mask = np.isfinite(time_data)
                time_data = time_data[valid_mask]
            else:
                valid_mask = np.all(np.isfinite(time_data), axis=0)
                time_data = time_data[:, valid_mask]
        
        # 创建清理后的时域数据
        cleaned_time_series = TimeSeriesData(
            data=time_data,
            sampling_rate=brain_wave_data.time_series.sampling_rate,
            timestamp=brain_wave_data.time_series.timestamp,
            channel_names=brain_wave_data.time_series.channel_names
        )
        
        # 返回清理后的数据
        return BrainWaveData(
            time_series=cleaned_time_series,
            frequency_data=brain_wave_data.frequency_data,
            session_id=brain_wave_data.session_id,
            processing_stage=brain_wave_data.processing_stage,
            metadata=brain_wave_data.metadata
        )
    
    @staticmethod
    def validate_and_clean_data(data: np.ndarray, threshold: float = 3.0) -> np.ndarray:
        """
        验证和清理数据，移除异常值
        
        Args:
            data: 输入数据数组
            threshold: 异常值检测阈值（标准差倍数）
            
        Returns:
            np.ndarray: 清理后的数据
        """
        if len(data) == 0:
            return data
        
        # 计算统计信息
        mean = np.mean(data)
        std = np.std(data)
        
        # 移除超出阈值的异常值
        mask = np.abs(data - mean) <= threshold * std
        cleaned_data = data[mask]
        
        return cleaned_data if len(cleaned_data) > 0 else data
    
    @staticmethod
    def create_filtered_data_pair(
        raw_brain_wave_data: 'BrainWaveData',
        filter_function=None
    ) -> tuple:
        """
        创建原始数据和过滤数据的配对
        
        Args:
            raw_brain_wave_data: 原始脑波数据
            filter_function: 过滤函数，如果为None则使用默认过滤
            
        Returns:
            tuple: (原始数据, 过滤数据)
        """
        # 原始数据保持不变
        raw_data = raw_brain_wave_data
        
        # 创建过滤后的数据副本
        filtered_time_series = DataUtils.create_time_series_data(
            data=raw_data.time_series.data,  # 这里可以应用滤波算法
            sampling_rate=raw_data.time_series.sampling_rate,
            timestamp=raw_data.time_series.timestamp
        )
        
        # 创建过滤后的频域数据
        filtered_frequency_data = DataUtils.create_frequency_data(
            delta=raw_data.frequency_data.delta,
            theta=raw_data.frequency_data.theta,
            low_alpha=raw_data.frequency_data.low_alpha,
            high_alpha=raw_data.frequency_data.high_alpha,
            low_beta=raw_data.frequency_data.low_beta,
            high_beta=raw_data.frequency_data.high_beta,
            low_gamma=raw_data.frequency_data.low_gamma,
            mid_gamma=raw_data.frequency_data.mid_gamma,
            attention=raw_data.frequency_data.attention,
            meditation=raw_data.frequency_data.meditation,
            poor_signal=raw_data.frequency_data.poor_signal,
            blink_strength=raw_data.frequency_data.blink_strength,
            raw_value=raw_data.frequency_data.raw_value,
            is_filtered=True,  # 标记为过滤数据
            timestamp=raw_data.frequency_data.timestamp
        )
        
        # 创建过滤后的完整数据对象
        filtered_data = DataUtils.create_brain_wave_data(
            time_series=filtered_time_series,
            frequency_data=filtered_frequency_data,
            session_id=raw_data.session_id,
            processing_stage="filtered",
            metadata={
                **raw_data.metadata,
                'is_filtered': True,
                'filter_applied': True
            }
        )
        
        return raw_data, filtered_data
    
    @staticmethod
    def validate_neuropy_data_quality(
        poor_signal: int,
        attention: float,
        meditation: float,
        min_quality_threshold: int = 100
    ) -> dict:
        """
        验证NeuroPy3数据质量
        
        Args:
            poor_signal: 信号质量值 (0-200)
            attention: 注意力指数 (0-100)
            meditation: 冥想指数 (0-100)
            min_quality_threshold: 最小质量阈值
            
        Returns:
            dict: 包含验证结果的字典
        """
        validation_result = {
            'is_valid': True,
            'quality_level': 'unknown',
            'issues': [],
            'recommendations': []
        }
        
        # 检查信号质量
        if poor_signal >= 200:
            validation_result['is_valid'] = False
            validation_result['quality_level'] = 'no_contact'
            validation_result['issues'].append('设备未正确接触')
            validation_result['recommendations'].append('请检查设备连接和佩戴位置')
        elif poor_signal > min_quality_threshold:
            validation_result['quality_level'] = 'poor'
            validation_result['issues'].append('信号质量较差')
            validation_result['recommendations'].append('请调整设备位置或清洁电极')
        elif poor_signal <= 25:
            validation_result['quality_level'] = 'excellent'
        elif poor_signal <= 50:
            validation_result['quality_level'] = 'good'
        else:
            validation_result['quality_level'] = 'fair'
        
        # 检查注意力和冥想指数范围
        if not (0 <= attention <= 100):
            validation_result['issues'].append('注意力指数超出有效范围')
        if not (0 <= meditation <= 100):
            validation_result['issues'].append('冥想指数超出有效范围')
        
        return validation_result
    
    @staticmethod
    def compare_filtered_unfiltered_data(
        raw_data: 'BrainWaveData',
        filtered_data: 'BrainWaveData'
    ) -> dict:
        """
        比较过滤前后的数据差异
        
        Args:
            raw_data: 原始数据
            filtered_data: 过滤后数据
            
        Returns:
            dict: 包含比较结果的字典
        """
        comparison = {
            'signal_quality_improvement': False,
            'frequency_changes': {},
            'attention_change': 0.0,
            'meditation_change': 0.0,
            'overall_improvement': False
        }
        
        # 比较信号质量
        raw_quality = raw_data.frequency_data.poor_signal
        filtered_quality = filtered_data.frequency_data.poor_signal
        comparison['signal_quality_improvement'] = filtered_quality < raw_quality
        
        # 比较频域数据
        frequency_bands = ['delta', 'theta', 'low_alpha', 'high_alpha', 
                          'low_beta', 'high_beta', 'low_gamma', 'mid_gamma']
        
        for band in frequency_bands:
            raw_value = getattr(raw_data.frequency_data, band)
            filtered_value = getattr(filtered_data.frequency_data, band)
            comparison['frequency_changes'][band] = {
                'raw': raw_value,
                'filtered': filtered_value,
                'change_percent': ((filtered_value - raw_value) / max(raw_value, 0.001)) * 100
            }
        
        # 比较注意力和冥想指数
        comparison['attention_change'] = (
            filtered_data.frequency_data.attention - raw_data.frequency_data.attention
        )
        comparison['meditation_change'] = (
            filtered_data.frequency_data.meditation - raw_data.frequency_data.meditation
        )
        
        # 评估整体改善情况
        improvements = [
            comparison['signal_quality_improvement'],
            abs(comparison['attention_change']) < 5,  # 变化不大说明稳定
            abs(comparison['meditation_change']) < 5
        ]
        comparison['overall_improvement'] = sum(improvements) >= 2
        
        return comparison